import asyncio
import json
import logging
from typing import Dict, Set
from fastapi import WebSocket
from app.schemas import PriceUpdateMessage, ContractUpdateMessage, BalanceUpdateMessage

logger = logging.getLogger(__name__)


class WebSocketManager:
    def __init__(self):
        # Store active connections
        self.active_connections: Set[WebSocket] = set()
        # Store user-specific connections
        self.user_connections: Dict[str, Set[WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, user_id: str = None):
        """Accept a WebSocket connection"""
        await websocket.accept()
        self.active_connections.add(websocket)
        
        if user_id:
            if user_id not in self.user_connections:
                self.user_connections[user_id] = set()
            self.user_connections[user_id].add(websocket)
        
        logger.info(f"WebSocket connected. Total connections: {len(self.active_connections)}")
    
    def disconnect(self, websocket: WebSocket, user_id: str = None):
        """Remove a WebSocket connection"""
        self.active_connections.discard(websocket)
        
        if user_id and user_id in self.user_connections:
            self.user_connections[user_id].discard(websocket)
            if not self.user_connections[user_id]:
                del self.user_connections[user_id]
        
        logger.info(f"WebSocket disconnected. Total connections: {len(self.active_connections)}")
    
    async def send_personal_message(self, message: dict, websocket: WebSocket):
        """Send a message to a specific WebSocket"""
        try:
            await websocket.send_text(json.dumps(message))
        except Exception as e:
            logger.error(f"Error sending personal message: {e}")
            self.active_connections.discard(websocket)
    
    async def send_to_user(self, message: dict, user_id: str):
        """Send a message to all connections of a specific user"""
        if user_id not in self.user_connections:
            return
        
        disconnected = set()
        for websocket in self.user_connections[user_id].copy():
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"Error sending message to user {user_id}: {e}")
                disconnected.add(websocket)
        
        # Clean up disconnected websockets
        for websocket in disconnected:
            self.disconnect(websocket, user_id)
    
    async def broadcast(self, message: dict):
        """Broadcast a message to all connected clients"""
        if not self.active_connections:
            return
        
        disconnected = set()
        for websocket in self.active_connections.copy():
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"Error broadcasting message: {e}")
                disconnected.add(websocket)
        
        # Clean up disconnected websockets
        for websocket in disconnected:
            self.disconnect(websocket)
    
    async def broadcast_price_update(self, symbol: str, price: float, timestamp: int):
        """Broadcast price update to all clients"""
        message = PriceUpdateMessage(
            symbol=symbol,
            price=price,
            timestamp=timestamp
        ).dict()
        
        await self.broadcast(message)
    
    async def send_contract_update(self, user_id: str, contract_data: dict):
        """Send contract update to specific user"""
        message = ContractUpdateMessage(
            contract=contract_data
        ).dict()
        
        await self.send_to_user(message, user_id)
    
    async def send_balance_update(self, user_id: str, balance: float):
        """Send balance update to specific user"""
        message = BalanceUpdateMessage(
            user_id=user_id,
            balance=balance
        ).dict()
        
        await self.send_to_user(message, user_id)


# Global WebSocket manager instance
websocket_manager = WebSocketManager()
